package com.dailyblue.project.estate.test;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.List;

/**
 * @Auther:Liu
 * @Testname:App
 * @Date:2025/9/23 14:33
 */
@SpringBootTest
@Slf4j
public class App {
    @Resource
    private VectorStore vectorStore;
    @Value("classpath:text-source.txt")
    private org.springframework.core.io.Resource textResource;
    @Value("classpath:pdf-source.pdf")
    private org.springframework.core.io.Resource pdfResource;
    @Value("classpath:markdown-source.md")
    private org.springframework.core.io.Resource markdownResource;

    /**
     * 添加数据
     */
    @Test
    public void A() {
        List<Document> documents = List.of(
                new Document("I like SpringAI"),
                new Document("I like SpringBoot"),
                new Document("I like ElasticSearch"),
                new Document("I like Spring Cloud")
        );
        vectorStore.add(documents);
    }

    /**
     * 删除数据
     */
    @Test
    public void remove() {
        List<String> ids = List.of("0a682e94-3f25-4397-abcf-9fb33b3632f5");
        vectorStore.delete(ids);
    }

    // 根据查询字符串进⾏相似性搜索，返回相似的⽂档列表。
    @Test
    public void find1() {
        // 默认情况下，返回的数据为4条。
        List<Document> documents = vectorStore.similaritySearch("I like");
        documents.forEach(e -> log.info("e:{}", e));
    }

    // 根据SearchRequest对象进⾏更复杂的相似性搜索。
    @Test
    public void find2() {
        String question = "哈利波特";
        int top = 6;
        double threshold = 0.2;
        SearchRequest searchRequest = SearchRequest.builder()
                .query(question) // 问题
                .topK(top)  // 返回的条数
                .similarityThreshold(threshold)

                // 相似度阈值
                .build();
        List<Document> documents = vectorStore.similaritySearch(searchRequest);
        documents.forEach(e -> log.info("e:{}", e));
    }

    @Test
    public void textAdd() {
        TextReader txtReader = new TextReader(textResource);
        List<Document> documents = txtReader.get();
        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }

    @Test
    public void pdfAdd() {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(pdfResource);
        List<Document> documents = pdfReader.get();
        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }

    @Test
    public void markDownAdd() {
        MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeCodeBlock(false)
                .withIncludeBlockquote(false)
                .withAdditionalMetadata("filename", markdownResource.getFilename())
                .build();
        MarkdownDocumentReader markdownDocumentReader = new
                MarkdownDocumentReader(markdownResource, config);
        List<Document> documents = markdownDocumentReader.get();

        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }
}
